import torch


class ExtendedEmbedding(torch.nn.Module):
    def __init__(self, embedding0: torch.nn.Embedding, embedding1: torch.nn.Embedding):
        super().__init__()
        self.embedding0 = embedding0
        self.embedding1 = embedding1

    def forward(self, input_ids):
        weight = torch.cat([self.embedding0.weight, self.embedding1.weight], dim=0)
        return torch.nn.functional.embedding(input_ids, weight)


class ExtendedLinear(torch.nn.Module):
    def __init__(self, linear0: torch.nn.Linear, linear1: torch.nn.Linear):
        super().__init__()
        self.linear0 = linear0
        self.linear1 = linear1

    def forward(self, input_ids):
        weight = torch.cat([self.linear0.weight, self.linear1.weight], dim=0)
        return torch.nn.functional.linear(input_ids, weight)
